package org.chnxi.datastructures.tree;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;

import java.util.Objects;

public class AVLTreeDemo {

    public static void main(String[] args) {

        System.out.println("第一棵树，测试左旋转：");
        int arr[] = {4,3,6,5,7,8};

        //创建AVL TREE
        AVLTree<AVLTEntity> avlTree = new AVLTree<AVLTEntity>();
        for (int i=0;i<arr.length;i++){
            avlTree.add(new AVLTreeNode<AVLTEntity>(new AVLTEntity(arr[i] , "name-"+i)));
        }

        System.out.println("对当前的数进行中序遍历：");
        avlTree.printByInfixOrder();

        System.out.println("平衡处理后：");
        System.out.println("树的高度："+avlTree.getRoot().height());
        System.out.println("树的根节点："+avlTree.getRoot());
        System.out.println("左子树的高度："+avlTree.getRoot().leftHeight());
        System.out.println("右子树的高度："+avlTree.getRoot().rightHeight());

        System.out.println("==================================================");

        System.out.println("第二棵树，测试右旋转：");
        int arr2[] = {10,12,8,9,7,6};

        //创建AVL TREE
        AVLTree<AVLTEntity> avlTree2 = new AVLTree<AVLTEntity>();
        for (int i=0;i<arr2.length;i++){
            avlTree2.add(new AVLTreeNode<AVLTEntity>(new AVLTEntity(arr2[i] , "name-"+i)));
        }

        System.out.println("对当前的数进行中序遍历：");
        avlTree2.printByInfixOrder();

        System.out.println("平衡处理后：");
        System.out.println("树的高度："+avlTree2.getRoot().height());
        System.out.println("树的根节点："+avlTree2.getRoot());
        System.out.println("左子树的高度："+avlTree2.getRoot().leftHeight());
        System.out.println("右子树的高度："+avlTree2.getRoot().rightHeight());

        System.out.println("==================================================");

        System.out.println("第三棵树，测试双旋转：");
        int arr3[] = {10,11,7,6,8,9};

        //创建AVL TREE
        AVLTree<AVLTEntity> avlTree3 = new AVLTree<AVLTEntity>();
        for (int i=0;i<arr3.length;i++){
            avlTree3.add(new AVLTreeNode<AVLTEntity>(new AVLTEntity(arr3[i] , "name-"+i)));
        }

        System.out.println("对当前的数进行中序遍历：");
        avlTree3.printByInfixOrder();

        System.out.println("平衡处理后：");
        System.out.println("树的高度："+avlTree3.getRoot().height());
        System.out.println("树的根节点："+avlTree3.getRoot());
        System.out.println("左子树的高度："+avlTree3.getRoot().leftHeight());
        System.out.println("右子树的高度："+avlTree3.getRoot().rightHeight());

    }

}

class AVLTree<E extends Comparable<E>>{

    private AVLTreeNode<E> root;

    /**
     * 返回根节点
     * @return
     */
    public AVLTreeNode getRoot(){
        return root;
    }

    /**
     * 查找节点
     * @param e
     * @return
     */
    public AVLTreeNode<E> search(E e){
        if(root == null){
            return null;
        }
        return root.search(e);
    }

    /**
     * 查找指定节点的父节点
     * @param e
     * @return
     */
    public AVLTreeNode<E> searchParent(E e){
        if(root == null){
            return null;
        }
        return root.searchParent(e);
    }

    /**
     * 实现功能：
     *   1. 找到最小节点的值
     *   2. 删除该节点
     * @param node 传入的节点（当做二叉树排序节点的根节点）
     * @return 返回以node为根节点的二叉排序树的最小节点的值
     */
    public E delRightTreeMin(AVLTreeNode<E> node){

        AVLTreeNode<E> target = node;
        //循环查找左子节点，就会找到最小值
        while(target.left != null){
            target = target.left;
        }
        //删除最小节点
        delNode(target.e);
        return target.e;
    }

    /**
     * 删除节点
     * @param e 要删除节点的值
     */
    public void delNode(E e){
        if(root == null){
            return;
        }
        //先找到该节点
        AVLTreeNode<E> node = search(e);
        if(node == null){ //要删除的节点未找到
            return;
        }

        //目标节点是根，且没有左右子节点
        if(root.left == null && root.right == null){
            root = null;
            return;
        }

        //查找删除节点的父节点
        AVLTreeNode<E> parent = searchParent(e);

        //如果要删除的节点是叶子节点
        if(node.left == null && node.right == null){
            //判断node是parent的左子节点还是右子节点
            if(parent.left != null && parent.left.e.equals(e)){ //删除左子节点
                parent.left = null;
            }else if(parent.right != null && parent.right.e.equals(e)){ //删除右子节点
                parent.right = null;
            }
        }else if(node.left != null && node.right != null){ //删除的节点有两颗子树
            E mine = delRightTreeMin(node.right);
            node.e = mine;
        }else{ //剩余只有一个子树的情况
            if(node.left != null){
                if(parent != null){
                    if(parent.left.e.equals(e)){
                        parent.left = node.left;
                    }else{
                        parent.right = node.left;
                    }
                }else{
                    root = node.left;
                }
            }else{
                if(parent != null){
                    if(parent.left.e.equals(e)){
                        parent.left = node.right;
                    }else{
                        parent.right = node.right;
                    }
                }else{
                    root = node.right;
                }
            }

        }

    }

    /**
     * 添加子节点
     * @param node
     */
    public void add(AVLTreeNode<E> node){
        if(root == null){
            root = node;
        }else{
            root.add(node);
        }
    }

    /**
     * 使用中序遍历打印
     */
    public void printByInfixOrder(){
        if(root != null){
            root.printByInfixOrder();
        }
    }
}

/**
 * 平衡二叉树的节点
 * @param <E>
 */
class AVLTreeNode<E extends Comparable<E>> {

    E e;

    AVLTreeNode<E> left;

    AVLTreeNode<E> right;

    public AVLTreeNode(E e){
        this.e = e;
    }

    /**
     * 返回左子树的高度
     * @return
     */
    public int leftHeight(){
        if(left == null){
            return 0;
        }
        return this.left.height();
    }

    /**
     * 返回右子树的高度
     * @return
     */
    public int rightHeight(){
        if(right == null){
            return 0;
        }
        return this.right.height();
    }

    /**
     * 返回当前节点的高度，以该节点为根几点的子树高度
     * @return
     */
    public int height(){
        return Math.max(left == null ? 0 : left.height() , right == null ? 0 : right.height()) +1;
    }

    /**
     * 查找节点
     * @param e 希望找到的节点的值
     * @return 找到返回节点，否则返回null
     */
    public AVLTreeNode<E> search(E e){
        int diff = this.e.compareTo(e);
        if(diff == 0){ //返回当前节点
            return this;
        } else if(diff > 0){ //小于当前节点，从左子树查找
            if(this.left != null){
                return this.left.search(e);
            }
        }else{ //大于当前节点，从右子树查找
            if(this.right != null){
                return this.right.search(e);
            }
        }
        return null;
    }

    /**
     * 查找目标节点的父节点
     * @param e 要查找节点的值
     * @return 返回目标节点的父节点，否则返回null
     */
    public AVLTreeNode<E> searchParent(E e){
        if((this.left != null && this.left.e.equals(e) ) ||
                (this.right != null && this.right.e.equals(e))){
            return this;
        } else {
            int diff = this.e.compareTo(e);
            if(diff > 0 && this.left != null){ //从左子树进行查找
                return this.left.searchParent(e);
            }else if(diff <= 0 && this.right != null){ //从右子树递归查找
                return this.right.searchParent(e);
            }else{
                return null;
            }
        }
    }

    /**
     * 添加节点
     * @param node
     */
    public void add(AVLTreeNode<E> node){
        if(node == null){
            return;
        }
        int diff = this.e.compareTo(node.e);
        if(diff < 0){
            if(this.right == null){
                this.right = node;
            }else{
                this.right.add(node); //递归向右子树添加
            }
        }else{
            if(this.left == null){
                this.left = node;
            }else{
                this.left.add(node); //递归向左子树添加
            }
        }

        int leftHeight = this.leftHeight();
        int rightHeight = this.rightHeight();

        //当添加完一个节点后，如果（右子树的高度 - 左子树的高度）> 1，左旋转
        if(rightHeight - leftHeight > 1){
            //如果它的右子树的左子树高度大于它的右子树高度
            if(this.right != null && this.right.leftHeight() > this.right.rightHeight()){
                //先对右子树进行右旋
                this.right.rightRotate();
            }
            //再对当前节点发生左旋
            leftRotate();
        }else if(leftHeight - rightHeight > 1){
            //如果它的左子树的右子树高度大于它的左子树高度
            if(this.left != null && this.left.rightHeight() > this.left.leftHeight()){
                //先对左子树进行左旋
                this.left.leftRotate();
            }
            //再对当前节点进行右旋
            rightRotate();
        }

    }

    /**
     * 右旋转
     */
    private void rightRotate(){
        System.out.println("发生一次右旋转。。。");
        //* 创建一个新的节点（NewNode），新的节点的值等于当前根节点的值；
        //    - newNode.value = curNode.val;
        AVLTreeNode newNode = new AVLTreeNode(this.e);
        //* 把NewNode的左子树设置为当前节点的左子树；
        //    - newNode.right = curNode.right;
        newNode.right = this.right;
        //* 把newNode的右子树设置为当前节点的右子树的左子树；
        //    - newNode.left = curNode.left.right;
        newNode.left = this.left.right;
        //* 把当前节点的值换为右子树节点的值；
        //    - curNode.value = curNode.left.value;
        //    - curNode.left = curNode.left.left;
        this.e = this.left.e;
        this.left = this.left.left;
        //* 把当前节点的左子树设置为newNode节点
        //    - curNode.right = newNode;
        this.right = newNode;
    }

    /**
     * 发生左旋转
     */
    private void leftRotate(){
        System.out.println("发生一次左旋转。。。");
        // 1. 创建一个新的节点（NewNode），新的节点的值等于当前根节点的值；
        AVLTreeNode newNode = new AVLTreeNode(this.e);
        // 2. 把NewNode的左子树设置为当前节点的左子树；
        newNode.left = this.left;
        // 3. 把newNode的右子树设置为当前节点的右子树的左子树；
        newNode.right = this.right.left;
        // 4. 把当前节点的值换为右子树节点的值；
        this.e = this.right.e;
        // 5. 把当前节点的右子树换为右子树的右子树
        this.right = this.right.right;
        // 6.把当前节点的左子树设置为newNode节点
        this.left = newNode;
    }

    /**
     * 使用中序遍历打印
     */
    public void printByInfixOrder(){
        if(this.left != null){
            this.left.printByInfixOrder();
        }
        System.out.println(this);
        if(this.right != null){
            this.right.printByInfixOrder();
        }
    }

    @Override
    public String toString() {
        return "AVLTreeNode{e=" + e + "}";
    }
}

@Data
@AllArgsConstructor
@NoArgsConstructor
class AVLTEntity implements Comparable<AVLTEntity>{

    private int id;

    private String name;

    @Override
    public int compareTo(AVLTEntity o) {
        return this.id - o.id;
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        AVLTEntity bstEntity = (AVLTEntity) o;
        return id == bstEntity.id;
    }

    @Override
    public int hashCode() {
        return Objects.hash(id);
    }
}